//
// Ephi - simulation of magnetic fields and particles
// Copyright (C) 2007 Indrek Mandre <indrek(at)mare.ee>
// For more information please see http://www.mare.ee/indrek/ephi/
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>
#include <vector>

#include "ephi.hpp"

#define RADIUS 0.15
#define WR 0.035
#define SPACING 0.01
#define COILPOT 10000
#define CURRENT 200000
#define COIL_MEMORY_USAGE 1000
#define CAGE_RADIUS 0.32

#define PFIELD_RESOLUTION 0.01
#define PMAP_RESOLUTION 0.01
#if 1
  #define OCTREE_RESOLUTION 0.012
  #define PMAP_SMOOTH 0.03
#else
  #define OCTREE_RESOLUTION 0.005
  #define PMAP_SMOOTH 0.015
#endif

#define INJECTSPREAD 0.005
#define PARTICLES 45000
#define CHARGE 8e-8
#define BMUL 1e6

class WbPlasma : public Plasma
{
public:
  WbPlasma () : Plasma(TREE_48)
  {
    // load or generate the bicubic coil field interpolator
    if ( !cd.load ("wb3.cgen") )
      {
        {
          Ephi::debug (Ephi::DEBUG_APP, "Generating the coil bicubic interpolator and saving into wb3.cgen\n");
          CoilGen gen;
          gen.radius = RADIUS;
          gen.wr = WR;
          gen.memory_usage = COIL_MEMORY_USAGE;
          gen.do_efield = false;
          gen.do_cubic = true;
          gen.coil_matrix_n = 1;
          gen.set_space (SPACING);
          gen.geom_radial = (CAGE_RADIUS + 0.02) * 1.42;
          gen.geom_axial = (CAGE_RADIUS + 0.02) * 1.42;
          gen.generate ("wb3.cgen");
        }

        if ( !cd.load ("wb3.cgen") )
          {
            Ephi::debug (Ephi::DEBUG_CRITICAL, "Wb3: unable to read in the generated wb3.cgen!\n");
            exit(-1);
          }
      }

    make_polywell_cube (statics, RADIUS, WR, SPACING, CURRENT, 0, CoilDataFactory<StaticCICU>(cd));

    // fill in the electro-static structures
    // the positively charged coils
    prec_t rr = RADIUS + (2 * WR + SPACING) / PREC_SQRT2;
    pfield.setResolution (PFIELD_RESOLUTION);
    pfield.addTorus (vect3d(-rr, 0, 0), vect3d(1, 0, 0), RADIUS, WR, COILPOT);
    pfield.addTorus (vect3d(rr, 0, 0), vect3d(-1, 0, 0), RADIUS, WR, COILPOT);
    pfield.addTorus (vect3d(0, -rr, 0), vect3d(0, 1, 0), RADIUS, WR, COILPOT);
    pfield.addTorus (vect3d(0, rr, 0), vect3d(0, -1, 0), RADIUS, WR, COILPOT);
    pfield.addTorus (vect3d(0, 0, rr), vect3d(0, 0, -1), RADIUS, WR, COILPOT);
    pfield.addTorus (vect3d(0, 0, -rr), vect3d(0, 0, 1), RADIUS, WR, COILPOT);

    // and the surrounding cage that compresses the potential gradient
    //pf.set_no_cdist (true);
    prec_t ores = pfield.setResolution (CAGE_RADIUS / 10.0);
    std::vector<vect3d> planes;
    planes.push_back (vect3d(0, 1, 0));
    planes.push_back (vect3d(1, 0, -1).normal());
    planes.push_back (vect3d(0, -1, 1).normal());
    vect3d n1 = (planes[1] % planes[0]).normal();
    vect3d n2 = (planes[0] % planes[2]).normal();
    vect3d n3 = (planes[2] % planes[1]).normal();
    pfield.addSphereSegment (CAGE_RADIUS, n1, n2, n3, 0, true, 2);
    pfield.setResolution(ores);
    pfield.set_no_cdist (false);

    qmul = 0;
    pmap_box = CAGE_RADIUS + 0.01;
    pmap_resolution = PMAP_RESOLUTION;
    pmap_smooth = PMAP_SMOOTH;
    stepmag = 1;

    octree_c0 = vect3d(0, 0, 0);
    octree_c1 = vect3d(CAGE_RADIUS + 0.02, CAGE_RADIUS + 0.02, CAGE_RADIUS + 0.02);
    octree_resolution = OCTREE_RESOLUTION;
    octree_radius = CAGE_RADIUS + 0.01;

    bmul = BMUL;

    //dyn.setDTM (0.01, 5);
    //dyn.setDTM (0.01, 10);
    dyn.setDTM (0.01, 1);
    calibrate_type = INITIAL;

    freerun = -1;

    init();
  }

  void raze ()
  {
    prec_t tot_energy = 0;
    for ( size_t i = 0; i < dyn.size(); i++ )
      {
        const vect3d& pos = dyn.get_pos(i);
        if ( pos.magnitude() > CAGE_RADIUS * CAGE_RADIUS )
          {
            printf ("Razed particle %u that hit the cage\n", (unsigned int)i);
            dyn.erase (i);
            i--;
            continue;
          }

        if ( statics.inContact (pos) )
          {
            printf ("Razed particle %u that hit the coils\n", (unsigned int)i);
            dyn.erase (i);
            i--;
            continue;
          }
        tot_energy += ElectroDynamics::ms2eV (dyn.get_v(i).length(), ELECTRON_MASS);
        tot_energy += 10000 - dyn.getPotential (dyn.get_pos(i));
      }

    printf ("Mid potential = %g\n", dyn.getPotential(vect3d(0, 0, 0)));
    printf ("Total energy = %.3feV\n", tot_energy);
  }

  void store (const char *name)
  {
    dyn.save (name);
  }

  void load (const char *path)
  {
    assert (dyn.load (path));
    init();
  }

  void inject ()
  {
    const transf3d_pair *tfms = get_cube_polywell_transforms48 ();
    for ( size_t i = 0; i < PARTICLES; i++ )
      {
        vect3d pos(0.25, 0, 0);
        pos.x += prec_t_drand() * 0.02 - 0.005;
        pos.y += prec_t_drand() * 0.02 - 0.005;
        pos.z += prec_t_drand() * 0.02 - 0.005;
        prec_t pot = dyn.getPotential (pos);
        pot -= 1000;
        assert (pot > 0);
        dyn.inject_electron (pos, vect3d(-1, 0, 0), pot);
        dyn.step();

        if ( !(i % 1000) )
          {
            Screen screen(800, 800, 2 * CAGE_RADIUS, 2 * CAGE_RADIUS);
            screen.draw_axes();
            statics.draw (screen);
            for ( size_t j = 0; j < dyn.size(); j++ )
              {
                const vect3d& pos = dyn.get_pos(j);
                for ( size_t k = 0; k < 48; k++ )
                    screen.set_pixel (tfms[k].f * pos, Screen::RED);
              }
            char buf[128];
            sprintf (buf, "wb3_inject_%06d.png", (int)i);
            screen.write (buf);
            sprintf (buf, "wb3_inject_%06d.dyn", (int)i);
            store (buf);
          }
      }
  }

  #define ADJUST_STEPS 10000
  void adjust (size_t n)
  {
    if ( n < ADJUST_STEPS )
      {
        qmul = (CHARGE / (48 * PARTICLES)) / ELEMENTARY_CHARGE;
        qmul *= (prec_t)n / (prec_t)ADJUST_STEPS;
      }
    else
      {
        qmul = (CHARGE / (48 * PARTICLES)) / ELEMENTARY_CHARGE;
      }
  }

  int freerun;

  CoilData cd;
};

int main (int argc, char *argv[])
{
  Ephi::setDebugLevel (Ephi::DEBUG_APP);
  size_t n = 1;

  WbPlasma plasma;

#if 0
  plasma.load ("out/wb3_07062.dyn");
  //plasma.init();
  //plasma.render ("out/wb3_00000", true);
  //plasma.store ("out/wb3_00000.dyn");
  n = 7063;
#else
  plasma.inject();
  plasma.render ("out/wb3_00000", true);
  plasma.store ("out/wb3_00000.dyn");
#endif

  while (1)
    {
      plasma.adjust (n);
      plasma.step();
      plasma.raze();
      char name[128];
      sprintf (name, "out/wb3_%05d", (int)n);
      plasma.render (name, true);
      sprintf (name, "out/wb3_%05d.dyn", (int)n);
      plasma.store (name);
      n++;
    }
  
  return 0;
}

